# 将异步函数转换成异步上下文管理器
def asynccontextmanager(async_func):
    class AsyncContextManager:
        def __init__(self, *args, **kwargs):
            self.args = args
            self.kwargs = kwargs

        async def __aenter__(self):
            self.gen = async_func(self.args, **self.kwargs)
            return await self.gen.__anext__()

        async def __aexit__(self):
            try:
                await self.gen.__anext__()
            except StopAsyncIteration:
                pass

    return AsyncContextManager


class AsyncExitStack:
    def __init__(self):
        self.contexts = []  # 存储多个上下文

    async def enter_async_context(self, context):
        result = await context.__aenter__()
        self.contexts.append(context)
        return result

    async def __aenter__(self):
        return self

    # 退出时按相反顺序清理所有上下文
    # 参数：异常类型、异常值、异常追踪
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        for context in reversed(self.contexts):
            await context.__aexit__()


class SessionManager:
    def __init__(self, server_name):
        # 服务器的名称
        self.server_name = server_name
        # 运行状态 默认值为False
        self.is_running = False

    def run(self):
        class SessionContext:
            def __init__(self, manager):
                self.manager = manager  # 保存管理器的引用

            # 进入上下文时，启动会话
            async def __aenter__(self):
                print(f"Starting {self.manager.server_name} session ...")
                self.manager.is_running = True
                return self

            async def __aexit__(self):
                print(f"Stopping {self.manager.server_name} session ...")
                self.manager.is_running = False

        # 返回会话上下文的实例
        return SessionContext(self)


class FastMCP:
    def __init__(self, name):
        self.name = name
        # 创建会话管理器
        self.session_manager = SessionManager(name)
        # 创建工具列表
        self.tools = []

    def tool(self):
        # 装饰器函数，将工具函数加入工具列表中
        def decorator(func):
            self.tools.append(func)
            return func

        return decorator

    # 将MCP服务器转为HTTP ASGI应用
    def streamable_http_app(self):
        class HTTPApp:
            def __init__(self, mcp_instance):
                self.mcp = mcp_instance

            # ASGI调用接口
            async def __call__(self, scope, send):
                # 如果路径为/mcp的话，处理MCP协议请求
                if scope["path"] == "/mcp":
                    await self.handle_mcp_request(scope, send)
                else:
                    await send({"type": "http.response.start", "status": 404})
                    await send({"type": "http.response.body", "body": "Not Found"})

            async def handle_mcp_request(self, scope, send):
                print(f"handling mcp request for {self.mcp.name}")

        return HTTPApp(self)


# 定义Starlette应用
class Starlette:
    def __init__(self, routes=None, lifespan=None):
        self.routes = routes  # 保存路由
        self.lifespan = lifespan  # 保存生命周期管理器

    async def __call__(self, scope, send):
        if scope["type"] == "lifespan":
            await self.handle_lifespan(scope, send)
            return
        if scope["type"] == "http":
            await self.handle_http(scope, send)

    async def handle_http(self, scope, send):
        path = scope["path"]
        for route in self.routes:
            if path.startswith(route.path):  # /echo/mcp /echo
                sub_scope = scope.copy()
                sub_scope["path"] = path[len(route.path) :]  # 移除路由前缀 /mcp
                await route.app(sub_scope, send)
                return
        await send({"type": "http.response.start", "status": 404})
        await send({"type": "http.response.body", "body": b"Not Found"})

    async def handle_lifespan(self, scope, send):
        if scope["message"] == "startup":
            if self.lifespan:
                # 创建生命周期上下文
                self.lifespan_context = self.lifespan(self)
                await self.lifespan_context.__aenter__()  # 进入生命周期上下文
            await send({"type": "lifespan.startup.complete"})
        if scope["message"] == "shutdown":
            if hasattr(self, "lifespan_context"):
                await self.lifespan_context.__aexit__()
            await send({"type": "lifespan.shutdown.complete"})


class Mount:
    def __init__(self, path, app):
        self.path = path
        self.app = app


echo_mcp = FastMCP(name="Echo Server")
math_mcp = FastMCP(name="Math Server")


# 使用装饰器定义生命周期管理器
@asynccontextmanager
async def lifespan(app):
    async with AsyncExitStack() as stack:
        await stack.enter_async_context(echo_mcp.session_manager.run())
        await stack.enter_async_context(math_mcp.session_manager.run())
        yield


app = Starlette(
    routes=[
        Mount("/echo", echo_mcp.streamable_http_app()),
        Mount("/math", math_mcp.streamable_http_app()),
    ],
    lifespan=lifespan,
)


async def start_server():
    print(f"server startup sequence")

    async def send(message):
        print(f"send {message}")  # 打印发送消息

    # 1.创建startup事件
    lifespan_scope = {"type": "lifespan", "message": "startup"}
    # 2.Starlette接收到startup事件，调用lifespan管理器
    await app(lifespan_scope, send)
    # 3.lifespan管理器创建AsyncExitStack
    # 4.AsyncExitStack进入两个session_manager.run上下文
    # 5.两个session_manager就启动成功了
    # 6.yield执行，Starlette应用进入服务运行状态
    # 7.客户端发起HTTP请求
    print(f"Starlette服务器正在处理请求")
    http_scope = {"type": "http", "method": "POST", "path": "/echo/mcp"}
    await app(http_scope, send)
    lifespan_scope = {"type": "lifespan", "message": "shutdown"}
    await app(lifespan_scope, send)


import asyncio

asyncio.run(start_server())
